package cn.infrabase.token;

import cn.infrabase.DateTime;
import cn.infrabase.Errors;
import cn.infrabase.Gid;
import io.jsonwebtoken.*;
import io.jsonwebtoken.impl.DefaultClaims;
import io.jsonwebtoken.impl.DefaultHeader;
import io.jsonwebtoken.security.InvalidKeyException;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;
import io.jsonwebtoken.security.WeakKeyException;

import javax.annotation.Nullable;
import java.security.Key;
import java.time.temporal.ChronoUnit;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * JWT令牌工具类。
 *
 * @author Weimin Gao
 * @since 1.0.0
 **/
class JwtUtils {

	/**
	 * 生成令牌。
	 *
	 * @param algo    令牌签名算法。
	 * @param key     令牌签名密钥。
	 * @param efftime 令牌有效期（单位秒小于等于0时无到期时间永久有效）。
	 * @param tid     令牌标识。
	 * @param data    令牌附加数据。
	 * @return 令牌对象。
	 */
	static Token generate(String algo, @Nullable byte[] key, long efftime, @Nullable Gid tid, @Nullable Map<String, String> data) {
		try {
			SignatureAlgorithm algorithm = algorithm(algo);
			DateTime issued = DateTime.now();
			DateTime expiry = expiry(issued, efftime);
			String jti = jti(tid);
			Map<String, Object> claims = new HashMap<>();
			JwtBuilder builder = Jwts.builder().setHeaderParam("typ", "JWT");
			builder.setId(jti);
			if (efftime > 0) {
				builder.setExpiration(expiry.toDate());
			}
			if (algorithm != SignatureAlgorithm.NONE) {
				builder.signWith(key(algorithm, key), algorithm);
			}
			claims.put("issued", issued.getValue());
			claims.put("efftime", efftime);
			claims.put("expiry", expiry.getValue());
			if (data != null) {
				claims.put("data", data);
			}
			builder.addClaims(claims);
			String token = builder.compact();
			return Token.builder()
				.type(TokenType.JWT.code())
				.algorithm(algo)
				.tid(jti)
				.issued(issued)
				.efftime(efftime)
				.expiry(expiry)
				.data(data)
				.value(token)
				.build();
		} catch (Exception cause) {
			throw generateError(cause);
		}
	}

	/**
	 * 解析令牌。
	 *
	 * @param token 令牌序号化为字符串的值。
	 * @param key   令牌签名密钥。
	 * @return 令牌对象。
	 */
	static Token parse(String token, byte[] key) {
		try {
			JwtParserBuilder parser = Jwts.parserBuilder();
			if (key != null) {
				parser.setSigningKey(key);
			}
			Jwt<?, ?> jwt = parser.build().parse(token);
			DefaultHeader<?> header = header(jwt);
			DefaultClaims claims = claims(jwt);
			assert header != null;
			assert claims != null;
			return Token.builder()
				.type(TokenType.JWT.code())
				.algorithm(algorithm(header))
				.tid(claims.getId())
				.efftime(efftime(claims))
				.issued(datetime(claims, "issued"))
				.expiry(datetime(claims, "expiry"))
				.data(data(claims))
				.value(token)
				.build();
		} catch (Exception cause) {
			throw parseError(cause);
		}
	}

	// -----------------------------------------------------------------------------------------------------------------

	private static DateTime datetime(DefaultClaims claims, String name) {
		String value = claims.get(name, String.class);
		if (value != null) {
			return DateTime.parse(value);
		}
		return null;
	}

	private static TokenException generateError(Exception cause) {
		// 生成令牌签名密码弱
		if (cause instanceof WeakKeyException) {
			return new TokenException(cause, Errors.GENERATE_TOKEN_SIGNATURE_KEY_WEAK);
		} else {
			if (cause instanceof InvalidKeyException) {
				return new TokenException(cause, Errors.GENERATE_TOKEN_SIGNATURE_KEY_INVALID);
			} else {
				return new TokenException(cause, Errors.GENERATE_TOKEN_ERROR);
			}
		}
	}

	private static TokenException parseError(Exception cause) {
		// 解析令牌格式不正确
		if (cause instanceof MalformedJwtException) {
			return new TokenException(cause, Errors.PARSE_TOKEN_MALFORMED);
		} else {
			// 解析令牌已过期
			if (cause instanceof ExpiredJwtException) {
				return new TokenException(cause, Errors.PARSE_TOKEN_EXPIRED);
			} else {
				// 解析令牌签名不匹配
				if (cause instanceof SignatureException) {
					return new TokenException(cause, Errors.PARSE_TOKEN_SIGNATURE_NOT_MATCH);
				} else {
					// 解析令牌发生错误
					return new TokenException(cause, Errors.PARSE_TOKEN_ERROR);
				}
			}
		}
	}

	@SuppressWarnings("unused")
	private static Key key(SignatureAlgorithm algo, byte[] key) {
		return Keys.hmacShaKeyFor(key);
	}

	private static String jti(Gid gid) {
		if (gid == null) {
			return Gid.get().value();
		} else {
			return gid.value();
		}
	}

	private static SignatureAlgorithm algorithm(String algo) {
		return SignatureAlgorithm.forName(algo);
	}

	private static long efftime(DefaultClaims claims) {
		Object value = claims.get("efftime");
		if (value == null) {
			return 0;
		} else {
			return Long.parseLong(value.toString());
		}
	}

	private static DateTime expiry(DateTime issued, long efftime) {
		if (efftime <= 0) {
			return DateTime.max();
		} else {
			return issued.plus(efftime, ChronoUnit.SECONDS);
		}
	}

	private static DefaultHeader<?> header(Jwt<?, ?> jwt) {
		Object header = jwt.getHeader();
		if (header instanceof DefaultHeader) {
			return (DefaultHeader<?>) header;
		}
		return null;
	}

	private static String algorithm(DefaultHeader<?> header) {
		Object algorithm = header.get("alg");
		return algorithm.toString();
	}

	private static DefaultClaims claims(Jwt<?, ?> jwt) {
		Object body = jwt.getBody();
		if (body instanceof DefaultClaims) {
			return (DefaultClaims) body;
		}
		return null;
	}

	private static Map<String, String> data(DefaultClaims claims) {
		Object object = claims.get("data");
		if (object != null) {
			if (object instanceof Map) {
				Map<String, String> data = new LinkedHashMap<>();
				Map<?, ?> values = (Map<?, ?>) object;
				values.forEach((key, value) -> data.put(String.valueOf(key), String.valueOf(value)));
				return data;
			}
		}
		return null;
	}
}
